import random
import os.path as osp
from glob import glob
import os


def write_rel_paths(phase, names, out_dir):
    """将文件相对路径存储在txt格式文件中"""
    with open(osp.join(out_dir, phase + '.txt'), 'w') as f:
        for name in names:
            f.write(name)
            f.write('\n')


def add_list(split):
    data_dir = osp.join('.', split, 'list')
    if not os.path.isdir(data_dir):
        os.makedirs(data_dir)
    names = sorted(list(map(osp.basename, glob(osp.join(split, 'A', '*.png')))))
    if split == 'train':
        train_ratio = 0.95
        len_train = int(len(names) * train_ratio)
        write_rel_paths('train', names[0:len_train], data_dir)
        write_rel_paths('val', names[len_train:], data_dir)
    elif split == 'test':
        write_rel_paths(split, names, data_dir)
    print(f'{split}数据集划分完成')


add_list('train')
add_list('test')
